# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： img_loader
Description :
Author : 'li'
date： 2022/1/27
-------------------------------------------------
Change Activity:
2022/1/27:
-------------------------------------------------
"""
import numpy as np

from ml.cv.misc.img_io import read_img


class LoadImage:
    def __init__(self, to_gray_img=False, to_float32=False, shape=None, input_key='path', output_key='ori_img'):
        """

        Args:
            to_gray_img: to gray img
            to_float32:
            shape: resize to image shape
        """
        self.to_gray_img = to_gray_img
        self.to_float32 = to_float32
        self.shape = shape
        self.input_key = input_key
        self.output_key = output_key

    def __call__(self, data):
        """

        Args:
            data:

        Returns:

        """
        path = data[self.input_key]
        img = read_img(path)
        if self.to_float32:
            img = img.astype(np.float32)
        data[self.output_key] = img
        return data
